import json
import re

from llm import Embedding, enc
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np


class SampleGenerator:
    def __init__(self, samples, func):
        self.samples = samples
        self.func = func
        self.embedding_matrix = []
        self.transformer_data()
        self.embedding_sample()

    def transformer_data(self):
        # 处理后应该有 embedding_content 和 return_content
        self.samples = [self.func(i) for i in self.samples]

    def embedding_sample(self):
        embedding_content_list = [i['embedding_content'] for i in self.samples]

        self.embedding_matrix = []

        batch_size = 64
        for i in range(0, len(embedding_content_list), batch_size):
            batch_content_list = embedding_content_list[i:i + batch_size]
            self.embedding_matrix.extend(Embedding(batch_content_list))

    def get_sample(self, content, top_n=3):
        embedding = Embedding(content)
        embedding = np.array(embedding).reshape(1, -1)
        similarities = cosine_similarity(embedding, self.embedding_matrix)
        top_indices = similarities.argsort()[0][-top_n:][::-1]
        return [self.samples[i]['return_content'] for i in top_indices[:top_n]]

    def add_sample(self, sample):
        sample = self.func(sample)
        self.samples.append(sample)
        self.embedding_matrix.append(Embedding(sample))


def func(x):
    return x


def get_exp_sample():
    li = []
    for path in ['submit/thought_exp.jsonl', 'submit/submit_exp_fu3.jsonl']:
        with open(path, encoding='utf8') as f:
            for i in f:
                data = json.loads(i)
                # print(data['exp'])
                if 'class ' not in data['exp']:
                    if '<thought>' in data['exp']:
                        li.append({"embedding_content": data['question'],
                                   "return_content": '<thought>' + data['exp'].split('<thought>', 1)[-1]})
                    else:
                        if len(enc.encode(data['exp'])) < 350:
                            li.append({"embedding_content": data['question'],
                                       "return_content": data['exp']})
    print('exp num:', len(li))
    return SampleGenerator(li, func)





from utils.utils import grouped_dict



def get_pos_doc(text):
    tmp = []
    li = []
    for k, v in grouped_dict.items():
        for p in v:
            if k in text and p['column_name'] in text:
                if p['column_description'] not in tmp:
                    li.append({'column_name_en': p['column_name'], "column_name_zh": p['column_description']})
                    tmp.append(p['column_description'])
    return li


def get_column_sample():
    li = []
    for path in ['submit/thought_exp.jsonl', 'submit/submit_exp_fu3.jsonl']:
        with open(path, encoding='utf8') as f:
            for i in f:
                data = json.loads(i)

                if 'class ' not in data['exp']:
                    if '```md' in data['exp']:
                        li.append({"embedding_content": data['question'],
                                   "return_content": {
                                       "question": data['question'],
                                       "columns": get_pos_doc(data['exp'].split('```md', 1)[-1])}
                                   })
    print('column sample num:', len(li))
    return SampleGenerator(li, func)


import pandas as pd
from collections import defaultdict
def get_sql_sample():
    li = []
    for path in ['submit/loga.jsonl','submit/logb.jsonl','submit/logc.jsonl']:
        group_data = []
        with open(path, encoding='utf8') as f:
            for i in f:
                data = json.loads(i)
                group_data.append(data)

        tid = []
        for i in group_data:
            tid.append(i['tid'])
        d = defaultdict(list)
        n=0
        for i in tid:
            for j in group_data:
                if j['tid'] == i:
                    if re.search('select',str(j['final_code'][0]),re.IGNORECASE):
                        d[i].append(j)
                        n+=1

        for i in d:
            content = ''
            for j in d[i]:
                content+=f"问题:{j['question']}\nSQL:{j['final_code'][0]}\n"
            for j in d[i]:
                li.append(
                    {"embedding_content": j['question'],
                     "return_content": content}
                )
        return SampleGenerator(li, func)
if __name__ == '__main__':
    get_sql_sample()
